Skip to content

TypeTree support in autodiff #144197

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 13 commits into
base: master
Choose a base branch
from
Open

TypeTree support in autodiff #144197

wants to merge 13 commits into from
+1,169 −13

Conversation

KMJ-007
Copy link
Contributor

@KMJ-007 KMJ-007 commented Jul 19, 2025 •
edited
Loading

TypeTrees for Autodiff

What are TypeTrees?

Memory layout descriptors for Enzyme. Tell Enzyme exactly how types are structured in memory so it can compute derivatives efficiently.

Structure

TypeTree(Vec<Type>)

Type {
    offset: isize,  // byte offset (-1 = everywhere)
    size: usize,    // size in bytes
    kind: Kind,     // Float, Integer, Pointer, etc.
    child: TypeTree // nested structure
}

Example: fn compute(x: &f32, data: &[f32]) -> f32

Input 0: x: &f32

TypeTree(vec![Type {
    offset: -1, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: -1, size: 4, kind: Float,
        child: TypeTree::new()
    }])
}])

Input 1: data: &[f32]

TypeTree(vec![Type {
    offset: -1, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: -1, size: 4, kind: Float,  // -1 = all elements
        child: TypeTree::new()
    }])
}])

Output: f32

TypeTree(vec![Type {
    offset: -1, size: 4, kind: Float,
    child: TypeTree::new()
}])

Why Needed?

  • Enzyme can't deduce complex type layouts from LLVM IR
  • Prevents slow memory pattern analysis
  • Enables correct derivative computation for nested structures
  • Tells Enzyme which bytes are differentiable vs metadata

What Enzyme Does With This Information:

Without TypeTrees (current state):

; Enzyme sees generic LLVM IR:
define float @distance(i8* %p1, i8* %p2) {
; Has to guess what these pointers point to
; Slow analysis of all memory operations
; May miss optimization opportunities
}

With TypeTrees (our implementation):

define "enzyme_type"="{[]:Float@float}" float @distance(
    ptr "enzyme_type"="{[]:Pointer}" %p1, 
    ptr "enzyme_type"="{[]:Pointer}" %p2
) {
; Enzyme knows exact type layout
; Can generate efficient derivative code directly
}

TypeTrees - Offset and -1 Explained

Type Structure

Type {
    offset: isize, // WHERE this type starts
    size: usize,   // HOW BIG this type is
    kind: Kind,    // WHAT KIND of data (Float, Int, Pointer)
    child: TypeTree // WHAT'S INSIDE (for pointers/containers)
}

Offset Values

Regular Offset (0, 4, 8, etc.)

Specific byte position within a structure

struct Point {
    x: f32, // offset 0, size 4
    y: f32, // offset 4, size 4
    id: i32, // offset 8, size 4
}

TypeTree for &Point (internal representation):

TypeTree(vec![
    Type { offset: 0, size: 4, kind: Float },   // x at byte 0
    Type { offset: 4, size: 4, kind: Float },   // y at byte 4
    Type { offset: 8, size: 4, kind: Integer }  // id at byte 8
])

Generates LLVM:

"enzyme_type"="{[]:Float@float}"

Offset -1 (Special: "Everywhere")

Means "this pattern repeats for ALL elements"

Example 1: Array [f32; 100]

TypeTree(vec![Type {
    offset: -1, // ALL positions
    size: 4,    // each f32 is 4 bytes
    kind: Float, // every element is float
}])

Instead of listing 100 separate Types with offsets 0,4,8,12...396

Example 2: Slice &[i32]

// Pointer to slice data
TypeTree(vec![Type {
    offset: -1, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: -1, // ALL slice elements
        size: 4,    // each i32 is 4 bytes
        kind: Integer
    }])
}])

Example 3: Mixed Structure

struct Container {
    header: i64,        // offset 0
    data: [f32; 1000],  // offset 8, but elements use -1
}
TypeTree(vec![
    Type { offset: 0, size: 8, kind: Integer }, // header
    Type { offset: 8, size: 4000, kind: Pointer,
        child: TypeTree(vec![Type {
            offset: -1, size: 4, kind: Float // ALL array elements
        }])
    }
])

@rustbot rustbot added F-autodiff `#![feature(autodiff)]` S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. labels Jul 19, 2025
@KMJ-007 KMJ-007 mentioned this pull request Jul 19, 2025
Closed
@rust-log-analyzer

This comment has been minimized.

Sign in to view
@rustbot rustbot added the A-LLVM Area: Code generation parts specific to LLVM. Both correctness bugs and optimization-related issues. label Jul 19, 2025
@rust-log-analyzer

This comment has been minimized.

Sign in to view
@KMJ-007
Copy link
Contributor Author

KMJ-007 commented Jul 19, 2025

Currently, I have implemented only for memcpy

@KMJ-007
Copy link
Contributor Author

KMJ-007 commented Jul 19, 2025

r? @ZuseZ4

@KMJ-007 KMJ-007 marked this pull request as ready for review July 19, 2025 23:50
@rustbot rustbot added S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. and removed S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Jul 19, 2025
@rustbot
Copy link
Collaborator

rustbot commented Jul 19, 2025

Some changes occurred in compiler/rustc_ast/src/expand/autodiff_attrs.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_codegen_llvm/src/builder/autodiff.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_codegen_ssa

cc @WaffleLapkin

Some changes occurred in compiler/rustc_monomorphize/src/partitioning/autodiff.rs

cc @ZuseZ4

@rust-log-analyzer

This comment has been minimized.

Sign in to view
@rust-log-analyzer

This comment has been minimized.

Sign in to view
@rust-log-analyzer

This comment has been minimized.

Sign in to view
@rust-log-analyzer

This comment has been minimized.

Sign in to view
@rust-log-analyzer

This comment has been minimized.

Sign in to view
@rustbot
Copy link
Collaborator

rustbot commented Jul 20, 2025

Some changes occurred in compiler/rustc_codegen_gcc

cc @antoyo, @GuillaumeGomez

@rust-log-analyzer

This comment has been minimized.

Sign in to view
@KMJ-007
Copy link
Contributor Author

KMJ-007 commented Jul 21, 2025

CI is failing, fixing them!

antoyo
antoyo reviewed Jul 21, 2025
View reviewed changes
@rustbot
Copy link
Collaborator

rustbot commented Jul 23, 2025

Some changes occurred in src/tools/enzyme

cc @ZuseZ4

@rustbot

This comment has been minimized.

Sign in to view
@rustbot rustbot added has-merge-commits PR has merge commits, merge with caution. S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Jul 23, 2025
@rustbot

This comment has been minimized.

Sign in to view
@rustbot rustbot removed the has-merge-commits PR has merge commits, merge with caution. label Jul 23, 2025
@rust-log-analyzer

This comment has been minimized.

Sign in to view
@rust-log-analyzer

This comment has been minimized.

Sign in to view
@rust-log-analyzer

This comment has been minimized.

Sign in to view
ZuseZ4
ZuseZ4 reviewed Aug 26, 2025
View reviewed changes
@rust-log-analyzer

This comment has been minimized.

Sign in to view
@rust-log-analyzer

This comment has been minimized.

Sign in to view
klensy
klensy reviewed Aug 27, 2025
View reviewed changes
typos.toml Outdated
@@ -34,6 +34,7 @@ targetting = "targetting"
publically = "publically"
clonable = "clonable"
moreso = "moreso"
EnzymeTypeTreeShiftIndiciesEq = "EnzymeTypeTreeShiftIndiciesEq"
Copy link
Contributor

@klensy klensy Aug 27, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

https://rust-lang.zulipchat.com/#narrow/channel/182449-t-compiler.2Fhelp/topic/.E2.9C.94.20tidy.20error.3A.20typo.20EnzymeTypeTreeShiftIndiciesEq.20to.20Enzym.2E.2E.2E/near/536350657: default.extend-identifiers block, not default.extend-words

@rustbot
Copy link
Collaborator

rustbot commented Aug 27, 2025

⚠️ Warning ⚠️

@KMJ-007 KMJ-007 requested a review from ZuseZ4 August 27, 2025 18:28
@rust-log-analyzer

This comment has been minimized.

Sign in to view
@ZuseZ4
Copy link
Member

ZuseZ4 commented Sep 2, 2025

Can you add your Example 3: Mixed Structure as a testcase? It would be great if you have a case demonstrating that the code doesn't just generate 1000 entries for float offsets, but optimizes it to use -1 since all cases are equal

struct Container {

	header: i64, // offset 0

	data: [f32; 1000], // offset 8, but elements use -1

}
TypeTree(vec![

	Type { offset: 0, size: 8, kind: Integer }, // header

	Type { offset: 8, size: 4000, kind: Pointer,

	child: TypeTree(vec![Type {

	offset: -1, size: 4, kind: Float // ALL array elements

}])

}

])

@KMJ-007 KMJ-007 mentioned this pull request Sep 4, 2025
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@antoyo antoyo antoyo left review comments

@ZuseZ4 ZuseZ4 Awaiting requested review from ZuseZ4

+1 more reviewer

@klensy klensy klensy left review comments

Reviewers whose approvals may not affect merge requirements
Assignees

@ZuseZ4 ZuseZ4

Labels
A-LLVM Area: Code generation parts specific to LLVM. Both correctness bugs and optimization-related issues. A-run-make Area: port run-make Makefiles to rmake.rs F-autodiff `#![feature(autodiff)]` S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
7 participants
@KMJ-007 @rust-log-analyzer @rustbot @bors @ZuseZ4 @antoyo @klensy